#include <arpa/inet.h>
#include "mod-tcp.hpp"


const int TcpDecoder::DEBUG = true;
const int TcpDecoder::DEBUG_TIMEOUT = (DEBUG && true);
const int TcpDecoder::DEBUG_DECODE  = (DEBUG && false);



Record * TcpDecoder::decode (Packet &pkt, Record &prev)
{
  TcpSession * ssn;
  TcpHdr * hdr = (TcpHdr*) pkt.getCur ();
  if (NULL == hdr)
    return NULL;

  const int hdrLen = hdr->thOff * 4;
  TcpRecord &r = this->rec;

  debug (DEBUG_DECODE, "Etner TCP decoder");

  // if incmpleted datagram, leave decoding
  if (NULL == hdr || ! pkt.moveCur (hdrLen))
    return NULL;

  // set basic information (header, session)
  r.setData (DEC_TCP, hdr, hdrLen); 
  r.sport = ntohs (hdr->thSrcPort);
  r.dport = ntohs (hdr->thDstPort);
  ssn = fetchTcpSession (prev, r.sport, r.dport);
  r.setSession (ssn);

  // get TCP segment length and data 
  if (0 < (r.segSize = pkt.getRemain ()))
    r.segData = static_cast<byte_t*> (pkt.getCur ());
  else
    r.segData = NULL;

#if 0
  if (r.segData)
    {
      printf ("[%d -> %d] ", r.sport, r.dport);
      for (int i = 0; i < r.segSize; i++)
	{
	  printf ("%c", (isgraph (r.segData [i]) ? r.segData [i] : '.'));
	}
      printf ("\n");
    }
#endif
  
  debug (DEBUG_DECODE, "Leave TCP decoder");
  return &(this->rec);
}


TcpSession * TcpDecoder::fetchTcpSession 
(Record &prev, u_int16_t &sport, u_int16_t &dport)
{
  TcpSessionRoot * root = NULL;
  TcpSession * ssn;
  void * ssnKey;

  if (prev.getSessionRootPtr ())
    {
      if (NULL == (root = (TcpSessionRoot*)prev.getSessionRoot ()))
	{
	  root = new TcpSessionRoot ();
	  bool ret = prev.setSessionRoot (root);
	  assert (ret);
	}
    }

  assert (root);

  ssnKey = this->makeSsnKey (&sport, &dport, 2);
  if (NULL == (ssn = dynamic_cast<TcpSession*> 
	       (root->ssn.lookup (ssnKey, 4))))
    {
      ssn = new TcpSession ();
      ssn->joinTable (root->ssn, ssnKey, 4);
    }

  ssn->setTimeout (this->getTimer (), SSN_TIMEOUT);

  return ssn;
}



void TcpDecoder::timeoutSession (void * obj)
{
  TcpSession * ssn = static_cast <TcpSession*> (obj);
  PTree<TcpSession> * tree = ssn->tree;
  ssn->tnode = NULL;
  tree->remove (ssn->getKey (), ssn->getKeyLen ());
}

NDecFilter * TcpDecoder::makeFilter (string &name, string &arg)
{
  TcpFilter * f = NULL;

  if (name == "tcp.src_port")
    f = new TcpFilter (TcpFilter::FT_SRC_PORT);
  else if (name == "tcp.dst_port")
    f = new TcpFilter (TcpFilter::FT_DST_PORT);
  else if (name == "tcp.pcre")
    f = new TcpFilter (TcpFilter::FT_PCRE);
  else if (name == "tcp.flags")
    f = new TcpFilter (TcpFilter::FT_FLAGS);

  return f;
}

NDecMpxFilter * TcpDecoder::makeMpxFilter (string &name)
{
  TcpMpxFilter * f = NULL;

  if (name == "tcp.src_port")
    f = new TcpMpxFilter (TcpMpxFilter::FT_SRC_PORT);  
  else if (name == "tcp.dst_port")
    f = new TcpMpxFilter (TcpMpxFilter::FT_DST_PORT);  

  return f;  
}


//==========================================================

bool TcpFilter::setupPort (std::string &arg)
{
  char * e = NULL;
  unsigned int n;
  n = strtoul (arg.c_str (), &e, 0);
  if (*e != '\0' || n > 0xffff)
    {
      this->setErrMsg ("tcp.(src|dst)_port: expected digit \"" +  arg + "\"");
      return false;
    }

  this->port = static_cast <u_int16_t> (n);
	
  return true;
}


bool TcpFilter::setupFlags (std::string &arg)
{
  struct flaglist
  {
    char c;
    u_int8_t f;
  } flagList [] =
      {
	{'F', TH_FIN},
	{'S', TH_SYN},
	{'R', TH_RST},
	{'P', TH_PUSH},
	{'A', TH_ACK},
	{'U', TH_URG},
	{'E', TH_ECE},
	{'C', TH_CWR},
	{0, 0}
      };

  char * p, sep = '/';
  int n, i;

  /* flag data */
  p = const_cast<char*>(arg.c_str ());
  for (n = 0; n < 2; n++)
    {
      u_int8_t * f;
      f = (n == 0) ? &this->flagData : &this->flagMask;
      for (; *p != sep && *p != '\0'; p++)
	{
	  for (i = 0; flagList[i].c != 0; i++)
	    {
	      if (toupper (*p) == toupper (flagList[i].c))
		{
		  *f |= flagList[i].f;
		  break;
		}
	      
	      if (flagList[i].c == 0)
		{
		  this->setErrMsg (string ("Invalid Flag Char:") + p);
		  return false;
		}
	    }
	}

      if (*p == '\0')
	break;
      else if (*p == sep)
	p++;
    }

  if (*p != '\0')
    {
      this->setErrMsg (string ("Invalid Flag Char:") + p);
      return false;
    }

  if (this->flagMask == 0)
    this->flagMask = 0xff;;

  debug (DEBUG, "flagData = %02X, flagMask = %02X", 
	 this->flagData, this->flagMask);
	
  return true;
}

bool TcpFilter::setup (string &arg)
{
  bool ret = false;
  switch (this->type)
    {
    case FT_SRC_PORT:
    case FT_DST_PORT:
      ret = this->setupPort (arg);
      break;

    case FT_FLAGS:
      ret = this->setupFlags (arg);
      break;

    case FT_PCRE:
      {
	const char * e = NULL;
	int eOff;
	this->re = pcre_compile (arg.c_str (), PCRE_CASELESS|PCRE_MULTILINE, 
				 &e, &eOff, NULL);
	if (NULL == this->re)
	  {
	    this->setErrMsg ("tcp.pcre: PCRE compile error \"" +  arg + "\"");
	    return false;
	  }
      }
      break;

    default:
      assert (0);
    }

  return true;
}


void TcpFilter::check (NDecResult &res)
{
  u_int16_t * tgt = NULL;
  TcpRecord * rp = dynamic_cast<TcpRecord*> (res.getRecord (DEC_TCP));
  TcpHdr * hdr = static_cast<TcpHdr*> (rp->getPtr ());

  if (NULL == rp)
    return ;
  
  switch (this->type)
    {
    case FT_SRC_PORT: 
    case FT_DST_PORT: 
      tgt = (this->type == FT_SRC_PORT) ? &rp->sport : &rp->dport;
      debug (DEBUG, "compare port:%d tgt:%d", this->port, *tgt);

      if (this->port == 0 || *tgt == this->port)
	this->pushResult (tgt, 2);

      break;

    case FT_PCRE:      
      {
	const int ovsize = 30;
	int ov[ovsize];
	void * rPtr = NULL;
	int rSize = 0;
	int rc;

	try {
	  if (NULL == rp->segData)
	    throw false;

	  if (NULL == this->re)
	    {
	      rPtr = rp->segData;
	      rSize = rp->segSize;
	      throw true;
	    }

	  if (0 >= (rc = pcre_exec (this->re, NULL, (char*)rp->segData,
				    rp->segSize, 0, 0, ov, ovsize)))
	    throw false;

	  rSize = ov[1] - ov[0];
	  rPtr  = rp->segData + ov[0];
	  throw true;
	}
	catch (bool r)
	  {
	    if (r && NULL != rPtr)
	      this->pushResult (rPtr, rSize);
	  }
      }
      break;

    case FT_FLAGS:
      if ((hdr->thFlags & this->flagMask) == this->flagData)
	this->pushResult (&hdr->thFlags, sizeof (hdr->thFlags));

      break;

    default:
      assert (0);
    }
}


//==========================================================

bool TcpMpxFilter::setup (string &arg, void * p)
{
  char * e = NULL;
  unsigned int n;
  n = strtoul (arg.c_str (), &e, 0);
  if (*e != '\0' || n > 0xffff)
    {
      //      this->setErrMsg ("tcp.(src|dst)_port: expected digit \"" +  arg + "\"");
      return false;
    }

  u_int16_t port = static_cast <u_int16_t> (n);	
  NDecMpxRes * res = this->portTable.fetch (&port, sizeof (port));
  res->addItem (p);
  return true;
}


int TcpMpxFilter::exec (NDecResult &res)
{
  u_int16_t * tgt = NULL;
  NDecMpxRes * mr;
  TcpRecord * rp = dynamic_cast<TcpRecord*> (res.getRecord (DEC_TCP));
  if (NULL == rp)
    return false;
  
  tgt = (this->type == FT_SRC_PORT) ? &rp->sport : &rp->dport;
  if (NULL == (mr = this->portTable.lookup (tgt, sizeof (u_int16_t))))
    return 0;

  this->setResult (mr);
  return mr->getCount ();
}
